﻿namespace zijian666.AnyExtensions;
/// <summary>
/// 提供对 <see cref="CancellationTokenSource"/> 的扩展方法。
/// </summary>
public static class CancellationTokenExtensions
{
    /// <summary>
    /// 注册一个对象的释放操作到 <see cref="CancellationTokenSource"/> 的取消令牌中。
    /// </summary>
    /// <param name="cts">目标 <see cref="CancellationTokenSource"/>。</param>
    /// <param name="obj">需要释放的对象，可以是 <see cref="IDisposable"/> 或 <see cref="IAsyncDisposable"/>。</param>
    /// <exception cref="ArgumentNullException">当 <paramref name="cts"/> 为 null 时抛出。</exception>
    public static void RegisterDispose(this CancellationTokenSource cts, object obj)
    {
        if (cts is null)
        {
            throw new ArgumentNullException(nameof(cts));
        }

        if (obj is IDisposable d1)
        {
            cts.Token.Register(d1.Dispose);
        }
#if NET6_0_OR_GREATER
        else if (obj is IAsyncDisposable d2)
        {
            cts.Token.Register(() => d2.DisposeAsync().Wait());
        }
#endif
    }

    /// <summary>
    /// 获取 <see cref="CancellationTokenSource"/> 的取消令牌。
    /// </summary>
    /// <param name="cts">目标 <see cref="CancellationTokenSource"/>。</param>
    /// <returns>如果已请求取消，则返回一个已取消的 <see cref="CancellationToken"/>；否则返回 <see cref="CancellationTokenSource.Token"/>。</returns>
    public static CancellationToken GetToken(this CancellationTokenSource cts)
    {
        if (cts.IsCancellationRequested)
        {
            return new CancellationToken(true);
        }
        return cts.Token;
    }

    /// <summary>
    /// 如果 <see cref="CancellationTokenSource"/> 已请求取消，则抛出 <see cref="OperationCanceledException"/>。
    /// </summary>
    /// <param name="cts">目标 <see cref="CancellationTokenSource"/>。</param>
    /// <param name="message">异常消息，可选。</param>
    /// <exception cref="OperationCanceledException">当 <see cref="CancellationTokenSource"/> 已请求取消时抛出。</exception>
    public static void ThrowIfCancellationRequested(this CancellationTokenSource cts, string? message = null)
    {
        if (cts.IsCancellationRequested)
        {
            throw new OperationCanceledException(message);
        }
    }
}
